﻿using System;
using System.IO;
using System.Collections.Generic;
using System.ComponentModel;
using System.Data;
using System.Drawing;
using System.Text;
using System.Windows.Forms;
using System.Security;
using System.Runtime.InteropServices;
using System.Security.Cryptography;

namespace PLock {
	partial class PasswordDlg : Form {
		private PasswordDlg(string title) {
			InitializeComponent();
			if (title != null)
				this.Text = title;
		}

		// I'm aware the handling of this data isn't particularly secure since
		// it's copied to managed memory at various times, but I guess this is
		// better than nothing
		private static SecureString _pass = null;

		public const string MasterDigestFile = "master.digest";
		public const int SaltSize = 64;
		public const int PBKDFIterations = 1500;

		private static bool _prompting = false;

		public static bool PromptForPassword() {
			return PromptForPassword(null);
		}
		public static bool PromptForPassword(string prompt) {
			if (_prompting)
				return false;
			if (_pass != null)
				return true;
			using (PasswordDlg dlg = new PasswordDlg(prompt)) {
				try {
					_prompting = true;
					if (dlg.ShowDialog() == DialogResult.OK) {
						_pass = new SecureString();
						foreach (char c in dlg.textBoxPass.Text)
							_pass.AppendChar(c);
						return true;
					}
				} finally {
					_prompting = false;
				}
			}
			return false;
		}

		public static void ForgetPassword() {
			if (_pass != null) {
				_pass.Dispose();
				_pass = null;
			}
		}

		public static bool HasPassword {
			get { return _pass != null; }
		}

		public static SelfDestruct GetKey(string filename) {
			return GetKey(filename, null, null);
		}
		public static SelfDestruct GetKey(string filename, string temppass, byte[] tempsalt) {
			if (temppass != null) {
				return new SelfDestruct(PBKDF2.GetBytes(MasterKey(tempsalt, temppass),
										Encoding.UTF8.GetBytes(filename),
										PBKDFIterations, SymmetricAlgorithm().KeySize / 8));
			} else {
				if (_pass == null)
					throw new Exception("Bug: should have prompted for password");

				IntPtr ptr = IntPtr.Zero;
				try {
					ptr = Marshal.SecureStringToCoTaskMemAnsi(_pass);
					using (SelfDestruct bytes = new SelfDestruct(new byte[_pass.Length])) {
						for (int i = 0; i < bytes.Bytes.Length; i++)
							bytes.Bytes[i] = Marshal.ReadByte(ptr, i);
						return new SelfDestruct(PBKDF2.GetBytes(MasterKey(LoadSalt(), bytes.Bytes),
												Encoding.UTF8.GetBytes(filename),
												PBKDFIterations, SymmetricAlgorithm().KeySize / 8));
					}
				} finally {
					if (ptr != IntPtr.Zero)
						Marshal.ZeroFreeCoTaskMemAnsi(ptr);
				}
			}
		}

		private void PasswordDlg_Load(object sender, EventArgs e) {
			textBoxPass.Select();
		}

		private void textBoxPass_TextChanged(object sender, EventArgs e) {
			buttonOK.Enabled = IsValidPassword(textBoxPass.Text);
		}

		public static bool IsValidPassword(string pass) {
			if (pass == "") return false;
			foreach (char c in pass) {
				int i = (int)c;
				if (i < 33 || i > 126)
					return false;
			}
			return true;
		}

		public static byte[] NewSalt() {
			RNGCryptoServiceProvider rng = new RNGCryptoServiceProvider();
			byte[] salt = new byte[SaltSize];
			rng.GetBytes(salt);
			return salt;
		}

		public static void SaveDigest(string pass, byte[] salt) {
			byte[] digest = ComputeDigest(salt, pass);

			using (FileStream fs = File.Create(Settings.Instance.FullPath(MasterDigestFile))) {
				fs.Write(new byte[] { (byte)'P', (byte)'L', (byte)'K', 0x02 }, 0, 4);
				fs.Write(salt, 0, salt.Length);
				fs.Write(digest, 0, digest.Length);
			}
		}

		private static byte[] MasterKey(byte[] salt, string pass) {
			return PBKDF2.GetBytes(pass, salt, PBKDFIterations, SymmetricAlgorithm().KeySize / 8);
		}
		private static byte[] MasterKey(byte[] salt, byte[] pass) {
			return PBKDF2.GetBytes(pass, salt, PBKDFIterations, SymmetricAlgorithm().KeySize / 8);
		}

		private static byte[] ComputeDigest(byte[] salt, string pass) {
			return PBKDF2.GetBytes(MasterKey(salt, pass), new byte[] {70,197,209,174,29,138,179,72,60,176,148,78,
															6,18,81,219,6,30,203,100,110,30,130,245,200,
															0,140,86,231,210,64,111}, PBKDFIterations, SymmetricAlgorithm().KeySize / 8);
		}

		public static HashAlgorithm HashAlgorithm() {
			return SHA256.Create();
		}
		public static SymmetricAlgorithm SymmetricAlgorithm() {
			SymmetricAlgorithm sym = Rijndael.Create();
			sym.BlockSize = 128;
			sym.KeySize = 256;
			sym.Mode = CipherMode.CBC;
			sym.Padding = PaddingMode.PKCS7;
			return sym;
		}

		private static byte[] LoadSalt() {
			using (FileStream fs = File.OpenRead(Settings.Instance.FullPath(MasterDigestFile))) {
				byte[] sig = new byte[4];
				if (fs.Read(sig, 0, sig.Length) < 4 ||
					sig[0] != 'P' || sig[1] != 'L' || sig[2] != 'K' || sig[3] != 0x02)
					throw new Exception("Cannot read salt: Not a P-Lock digest file");
				byte[] salt = new byte[SaltSize];
				if (fs.Read(salt, 0, salt.Length) < salt.Length)
					throw new Exception("Cannot read salt: File content is invalid");
				return salt;
			}
		}

		private void buttonOK_Click(object sender, EventArgs e) {
			try {
				using (FileStream fs = File.OpenRead(Settings.Instance.FullPath(MasterDigestFile))) {
					byte[] sig = new byte[4];
					if (fs.Read(sig, 0, sig.Length) < 4 ||
						sig[0] != 'P' || sig[1] != 'L' || sig[2] != 'K' || sig[3] != 0x02)
						throw new Exception("Not a P-Lock digest file");
					byte[] salt = new byte[SaltSize];
					if (fs.Read(salt, 0, salt.Length) < salt.Length)
						throw new Exception("File content is invalid");
					byte[] computed = ComputeDigest(salt, textBoxPass.Text);
					byte[] disk = new byte[computed.Length];
					if (fs.Read(disk, 0, disk.Length) < disk.Length)
						throw new Exception("File content is invalid");
					for (int i = 0; i < computed.Length; i++) {
						if (disk[i] != computed[i]) {
							MessageBox.Show(this, "You have entered an incorrect password", "Incorrect Password", MessageBoxButtons.OK, MessageBoxIcon.Asterisk);
							textBoxPass.Text = "";
							return;
						}
					}
				}
			} catch (Exception ex) {
				MessageBox.Show(this, ex.Message, "Error reading digest file", MessageBoxButtons.OK, MessageBoxIcon.Error);
				return;
			}
			DialogResult = DialogResult.OK;
		}
	}

	class SelfDestruct : IDisposable {
		public SelfDestruct(byte[] bytes) {
			_bytes = bytes;
		}
		private byte[] _bytes;
		public byte[] Bytes {
			get { return _bytes; }
		}
		public void Dispose() {
			if (_bytes != null) {
				for (int i = 0; i < _bytes.Length; i++)
					_bytes[i] = 0;
				_bytes = null;
			}
		}
	}
}
